import pandas as pd
import re
import os

def clean_str(s):
    if pd.isna(s):
        return ''
    return re.sub(r'\s+', '', str(s)).strip()

def parse_colors_with_qty(color_str):
    """
    支持以下格式解析：
    - '黑15，白15' → {'黑': 15, '白': 15}
    - '黑 白各3' → {'黑': 3, '白': 3}
    - '黑、白各3' → {'黑': 3, '白': 3}
    - '黑/白各2' → {'黑': 2, '白': 2}
    """
    color_str = clean_str(color_str.replace("共", ""))

    # 特殊处理“各X”的情况
    match = re.search(r'各\s*(\d+)', color_str)
    if match:
        qty = int(match.group(1))
        # 获取“各”前的颜色部分
        color_part = color_str[:match.start()]
        colors = re.split(r'[、，,|/\\;；\s]+', color_part)
        return {color: qty for color in colors if color}

    # 默认按 “颜色+数字” 配对拆分
    pattern = r'([\u4e00-\u9fa5a-zA-Z]+)(\d+)'  # 如：杏15，白15
    matches = re.findall(pattern, color_str)
    return {color: int(qty) for color, qty in matches}

def normalize_color(color):
    """
    将颜色标准化，比如：白 → 白色，麻本 → 麻本色
    """
    if not color:
        return ""

    color = color.strip()
    
    # 例外：以下词结尾就不加"色"
    # exceptions = ['灰', '花', '彩', '黑', '蓝', '红', '绿', '金', '银', '紫', '粉']
    if color.endswith('色') or '色' in color: # or any(color.endswith(ex) for ex in exceptions)
        return color
    else:
        return color + '色'
    
def parse_sizes_and_qty(size_str):
    """
    处理码号字段：
    - 如果只包含数字 → 表示“件数”，返回 ['均码'], 件数
    - 如果包含 m, l, xl, xxl, s, 3xl 等字样 → 说明是尺码，返回尺码列表, None
    """
    size_str = clean_str(size_str)
    if not size_str:
        return ['均码'], None

    size_str = size_str.lower()

    # 如果只含数字就是件数
    if re.fullmatch(r'\d+', size_str):
        return ['均码'], int(size_str)

    # 标准尺码映射
    size_map = {
        'xs': 'XS',
        's': 'S',
        'm': 'M',
        'l': 'L',
        'xl': 'XL',
        'xxl': 'XXL',
        '3xl': '3XL',
        '4xl': '4XL',
        '5xl': '5XL',
        '均码': '均码'
    }

    # 按顺序匹配最长的合法尺码标记
    pattern = r'(3xl|4xl|5xl|xxl|xl|xs|s|m|l|均码|\d{2,3})'
    matches = re.findall(pattern, size_str)

    sizes = [size_map.get(m, m.upper()) for m in matches]
    return sizes or ['均码'], None

def preprocess_data(input_file='Restocktable.xlsx', output_file='Preprocessed_Data.xlsx'):
    df = pd.read_excel(os.path.join(os.getcwd(), input_file), header=1)

    processed_data = []
    serial_number = 1

    for _, row in df.iterrows():
        if str(row.get('货号', '')).strip() == '':
            continue

        # 获取字段
        supplier = clean_str(row.get('供应商', ''))
        item_code = clean_str(row.get('货号', ''))
        raw_color = clean_str(row.get('颜色', ''))
        raw_qty = clean_str(row.get('件数', ''))
        raw_size = clean_str(row.get('码号', ''))
        region = clean_str(row.get('地区', ''))

        # Step 1: 判断颜色字段是否带数量
        color_qty_map = parse_colors_with_qty(raw_color)
        if not color_qty_map:
            # 没带数量，就简单切分颜色
            color_qty_map = {color: None for color in re.split(r'[，、,|/\\;；\s]+', raw_color) if color}

        # Step 2: 解析尺码和是否数量写在码号中
        sizes, size_qty = parse_sizes_and_qty(raw_size)

        for color, color_qty in color_qty_map.items():
            for size in sizes:
                # 优先用颜色里的数量
                final_qty = color_qty
                if final_qty is None:
                    if size_qty is not None:
                        final_qty = size_qty
                    else:
                        # 从“件数”列中找数字
                        match = re.search(r'\d+', raw_qty)
                        final_qty = int(match.group()) if match else 1

                processed_data.append({
                    '序号': serial_number,
                    '供应商': supplier,
                    '货号': item_code,
                    '颜色':  normalize_color(color),
                    '码号': size,
                    '件数': final_qty,
                    '地区': region
                })
                serial_number += 1

    result_df = pd.DataFrame(processed_data)
    output_file = os.path.join(os.getcwd(), 'Preprocessed_Data.xlsx')
    result_df.to_excel(output_file, index=False)
    print(f"✅ 预处理完成，已保存到: {output_file}")

if __name__ == "__main__":
    preprocess_data()
